"""Security Audit and Logging System.

Provides comprehensive security event logging, monitoring, and reporting
for BLE authentication and communication security.
"""

import asyncio
import gzip
import json
import logging
import os
import time
from typing import Dict, List, Optional, Set, Any, Callable
from enum import Enum
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta
from pathlib import Path
from collections import defaultdict, deque

logger = logging.getLogger(__name__)


# Audit Configuration
AUDIT_CONFIG = {
    "log_level": "INFO",
    "retention_days": 7,
    "max_file_size": 10 * 1024 * 1024,  # 10MB
    "compression": True,
    "rotation_interval": 3600,          # 1 hour
    "batch_size": 100,
    "flush_interval": 30,               # 30 seconds
    "sensitive_fields": ["pin", "key", "password", "secret", "token"],
    "log_directory": "logs/security"
}


class EventType(Enum):
    """Security event types."""
    AUTH_ATTEMPT = "auth_attempt"           # Authentication attempt
    AUTH_SUCCESS = "auth_success"           # Successful authentication
    AUTH_FAILURE = "auth_failure"           # Failed authentication
    PAIRING_REQUEST = "pairing_request"     # Device pairing request
    PAIRING_SUCCESS = "pairing_success"     # Successful pairing
    PAIRING_FAILURE = "pairing_failure"     # Failed pairing
    SESSION_CREATED = "session_created"     # Session created
    SESSION_EXPIRED = "session_expired"     # Session expired
    KEY_GENERATED = "key_generated"         # Key generated
    KEY_ROTATED = "key_rotated"             # Key rotated
    KEY_REVOKED = "key_revoked"             # Key revoked
    DEVICE_BLOCKED = "device_blocked"       # Device blocked
    DEVICE_UNBLOCKED = "device_unblocked"   # Device unblocked
    SECURITY_VIOLATION = "security_violation"  # Security violation
    REPLAY_ATTACK = "replay_attack"         # Replay attack detected
    INVALID_MESSAGE = "invalid_message"     # Invalid message received
    SYSTEM_START = "system_start"           # System started
    SYSTEM_STOP = "system_stop"             # System stopped


class Severity(Enum):
    """Event severity levels."""
    LOW = "low"           # Informational
    MEDIUM = "medium"     # Warning
    HIGH = "high"         # Error
    CRITICAL = "critical" # Critical security issue


@dataclass
class SecurityEvent:
    """Security audit event."""
    event_id: str
    event_type: EventType
    severity: Severity
    timestamp: float = field(default_factory=time.time)
    session_id: Optional[str] = None
    device_id: Optional[str] = None
    user_id: Optional[str] = None
    source_ip: Optional[str] = None
    details: Dict[str, Any] = field(default_factory=dict)
    tags: Set[str] = field(default_factory=set)
    
    @property
    def datetime(self) -> datetime:
        """Get event datetime."""
        return datetime.fromtimestamp(self.timestamp)
    
    def add_detail(self, key: str, value: Any):
        """Add event detail."""
        # Filter sensitive information
        if any(sensitive in key.lower() for sensitive in AUDIT_CONFIG["sensitive_fields"]):
            value = "[REDACTED]"
        self.details[key] = value
    
    def add_tag(self, tag: str):
        """Add event tag."""
        self.tags.add(tag)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "event_id": self.event_id,
            "event_type": self.event_type.value,
            "severity": self.severity.value,
            "timestamp": self.timestamp,
            "datetime": self.datetime.isoformat(),
            "session_id": self.session_id,
            "device_id": self.device_id,
            "user_id": self.user_id,
            "source_ip": self.source_ip,
            "details": self.details,
            "tags": list(self.tags)
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'SecurityEvent':
        """Create from dictionary."""
        event = cls(
            event_id=data["event_id"],
            event_type=EventType(data["event_type"]),
            severity=Severity(data["severity"]),
            timestamp=data["timestamp"]
        )
        
        event.session_id = data.get("session_id")
        event.device_id = data.get("device_id") 
        event.user_id = data.get("user_id")
        event.source_ip = data.get("source_ip")
        event.details = data.get("details", {})
        event.tags = set(data.get("tags", []))
        
        return event


class AuditLogger:
    """Security audit logger with rotation and compression."""
    
    def __init__(self, log_directory: str = None):
        self.log_directory = Path(log_directory or AUDIT_CONFIG["log_directory"])
        self.log_directory.mkdir(parents=True, exist_ok=True)
        
        self.current_file: Optional[Path] = None
        self.current_handle = None
        self.file_size = 0
        self.event_buffer: deque = deque()
        
        # Background tasks
        self.flush_task: Optional[asyncio.Task] = None
        self.rotation_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        logger.info(f"AuditLogger initialized, directory: {self.log_directory}")
    
    async def start(self):
        """Start audit logger."""
        if self.is_running:
            return
        
        self.is_running = True
        self._create_new_log_file()
        
        # Start background tasks
        self.flush_task = asyncio.create_task(self._flush_loop())
        self.rotation_task = asyncio.create_task(self._rotation_loop())
        
        logger.info("AuditLogger started")
    
    async def stop(self):
        """Stop audit logger."""
        self.is_running = False
        
        # Cancel background tasks
        if self.flush_task:
            self.flush_task.cancel()
            try:
                await self.flush_task
            except asyncio.CancelledError:
                pass
        
        if self.rotation_task:
            self.rotation_task.cancel()
            try:
                await self.rotation_task
            except asyncio.CancelledError:
                pass
        
        # Flush remaining events
        await self._flush_events()
        
        # Close current file
        if self.current_handle:
            self.current_handle.close()
            self.current_handle = None
        
        logger.info("AuditLogger stopped")
    
    def _create_new_log_file(self):
        """Create new log file."""
        if self.current_handle:
            self.current_handle.close()
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"security_audit_{timestamp}.log"
        self.current_file = self.log_directory / filename
        
        self.current_handle = open(self.current_file, 'w', encoding='utf-8')
        self.file_size = 0
        
        logger.info(f"Created new audit log file: {filename}")
    
    def log_event(self, event: SecurityEvent):
        """Log security event."""
        # Add to buffer
        self.event_buffer.append(event)
        
        # Flush if buffer is full
        if len(self.event_buffer) >= AUDIT_CONFIG["batch_size"]:
            asyncio.create_task(self._flush_events())
    
    async def _flush_events(self):
        """Flush buffered events to disk."""
        if not self.event_buffer or not self.current_handle:
            return
        
        try:
            while self.event_buffer:
                event = self.event_buffer.popleft()
                event_line = json.dumps(event.to_dict()) + '\n'
                
                self.current_handle.write(event_line)
                self.file_size += len(event_line.encode('utf-8'))
            
            self.current_handle.flush()
            
        except Exception as e:
            logger.error(f"Error flushing audit events: {e}")
    
    async def _flush_loop(self):
        """Background flush loop."""
        while self.is_running:
            try:
                await asyncio.sleep(AUDIT_CONFIG["flush_interval"])
                await self._flush_events()
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Flush loop error: {e}")
    
    async def _rotation_loop(self):
        """Background log rotation loop."""
        while self.is_running:
            try:
                await asyncio.sleep(AUDIT_CONFIG["rotation_interval"])
                
                # Check if rotation is needed
                if (self.file_size >= AUDIT_CONFIG["max_file_size"] or
                    self._should_rotate_by_time()):
                    
                    await self._rotate_log_file()
                
                # Clean up old files
                await self._cleanup_old_files()
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Rotation loop error: {e}")
    
    def _should_rotate_by_time(self) -> bool:
        """Check if log should rotate by time."""
        if not self.current_file:
            return False
        
        file_age = time.time() - self.current_file.stat().st_mtime
        return file_age >= AUDIT_CONFIG["rotation_interval"]
    
    async def _rotate_log_file(self):
        """Rotate current log file."""
        if not self.current_file:
            return
        
        # Flush remaining events
        await self._flush_events()
        
        # Compress old file if enabled
        if AUDIT_CONFIG["compression"]:
            await self._compress_file(self.current_file)
        
        # Create new file
        self._create_new_log_file()
        
        logger.info("Log file rotated")
    
    async def _compress_file(self, file_path: Path):
        """Compress log file."""
        try:
            compressed_path = file_path.with_suffix(file_path.suffix + '.gz')
            
            with open(file_path, 'rb') as f_in:
                with gzip.open(compressed_path, 'wb') as f_out:
                    f_out.writelines(f_in)
            
            # Remove original file
            file_path.unlink()
            
            logger.debug(f"Compressed log file: {compressed_path}")
            
        except Exception as e:
            logger.error(f"Error compressing file {file_path}: {e}")
    
    async def _cleanup_old_files(self):
        """Clean up old log files."""
        try:
            cutoff_time = time.time() - (AUDIT_CONFIG["retention_days"] * 24 * 3600)
            
            for file_path in self.log_directory.glob("security_audit_*.log*"):
                if file_path.stat().st_mtime < cutoff_time:
                    file_path.unlink()
                    logger.debug(f"Deleted old audit file: {file_path}")
            
        except Exception as e:
            logger.error(f"Error cleaning up old files: {e}")
    
    async def read_events(self, start_time: float = None, 
                         end_time: float = None,
                         event_types: List[EventType] = None,
                         device_id: str = None,
                         session_id: str = None) -> List[SecurityEvent]:
        """Read events from log files."""
        events = []
        
        try:
            # Find relevant log files
            log_files = list(self.log_directory.glob("security_audit_*.log*"))
            log_files.sort()
            
            for file_path in log_files:
                file_events = await self._read_events_from_file(
                    file_path, start_time, end_time, event_types, device_id, session_id
                )
                events.extend(file_events)
            
            # Sort by timestamp
            events.sort(key=lambda e: e.timestamp)
            
        except Exception as e:
            logger.error(f"Error reading audit events: {e}")
        
        return events
    
    async def _read_events_from_file(self, file_path: Path,
                                   start_time: float = None,
                                   end_time: float = None,
                                   event_types: List[EventType] = None,
                                   device_id: str = None,
                                   session_id: str = None) -> List[SecurityEvent]:
        """Read events from a single log file."""
        events = []
        
        try:
            # Handle compressed files
            if file_path.suffix == '.gz':
                file_handle = gzip.open(file_path, 'rt', encoding='utf-8')
            else:
                file_handle = open(file_path, 'r', encoding='utf-8')
            
            with file_handle as f:
                for line in f:
                    try:
                        event_data = json.loads(line.strip())
                        event = SecurityEvent.from_dict(event_data)
                        
                        # Apply filters
                        if start_time and event.timestamp < start_time:
                            continue
                        if end_time and event.timestamp > end_time:
                            continue
                        if event_types and event.event_type not in event_types:
                            continue
                        if device_id and event.device_id != device_id:
                            continue
                        if session_id and event.session_id != session_id:
                            continue
                        
                        events.append(event)
                        
                    except (json.JSONDecodeError, KeyError) as e:
                        logger.warning(f"Invalid audit log line: {e}")
                        continue
        
        except Exception as e:
            logger.error(f"Error reading file {file_path}: {e}")
        
        return events


class SecurityAudit:
    """Main security audit system."""
    
    def __init__(self, log_directory: str = None):
        self.audit_logger = AuditLogger(log_directory)
        self.event_counters: Dict[str, int] = defaultdict(int)
        self.alert_thresholds = {
            "failed_auth_per_hour": 5,
            "blocked_devices_per_hour": 3,
            "replay_attacks_per_hour": 2
        }
        
        # Alert callbacks
        self.alert_callbacks: List[Callable[[str, Dict], None]] = []
        
        logger.info("SecurityAudit initialized")
    
    async def start(self):
        """Start security audit system."""
        await self.audit_logger.start()
        
        # Log system start
        await self.log_system_event("system_start", Severity.LOW, {
            "component": "SecurityAudit",
            "version": "1.0.0"
        })
        
        logger.info("SecurityAudit started")
    
    async def stop(self):
        """Stop security audit system."""
        # Log system stop
        await self.log_system_event("system_stop", Severity.LOW, {
            "component": "SecurityAudit"
        })
        
        await self.audit_logger.stop()
        logger.info("SecurityAudit stopped")
    
    def add_alert_callback(self, callback: Callable[[str, Dict], None]):
        """Add security alert callback."""
        self.alert_callbacks.append(callback)
    
    def _generate_event_id(self) -> str:
        """Generate unique event ID."""
        import secrets
        return f"evt_{int(time.time())}_{secrets.token_urlsafe(8)}"
    
    def _trigger_alerts(self, event: SecurityEvent):
        """Check and trigger security alerts."""
        current_hour = int(time.time() / 3600)
        
        # Count events by type and hour
        if event.event_type == EventType.AUTH_FAILURE:
            key = f"failed_auth_{current_hour}"
            self.event_counters[key] += 1
            
            if self.event_counters[key] >= self.alert_thresholds["failed_auth_per_hour"]:
                self._send_alert("HIGH_AUTH_FAILURES", {
                    "count": self.event_counters[key],
                    "hour": current_hour,
                    "device_id": event.device_id
                })
        
        elif event.event_type == EventType.DEVICE_BLOCKED:
            key = f"blocked_devices_{current_hour}"
            self.event_counters[key] += 1
            
            if self.event_counters[key] >= self.alert_thresholds["blocked_devices_per_hour"]:
                self._send_alert("HIGH_DEVICE_BLOCKS", {
                    "count": self.event_counters[key],
                    "hour": current_hour
                })
        
        elif event.event_type == EventType.REPLAY_ATTACK:
            key = f"replay_attacks_{current_hour}"
            self.event_counters[key] += 1
            
            if self.event_counters[key] >= self.alert_thresholds["replay_attacks_per_hour"]:
                self._send_alert("REPLAY_ATTACK_WAVE", {
                    "count": self.event_counters[key],
                    "hour": current_hour,
                    "device_id": event.device_id
                })
    
    def _send_alert(self, alert_type: str, details: Dict[str, Any]):
        """Send security alert."""
        for callback in self.alert_callbacks:
            try:
                callback(alert_type, details)
            except Exception as e:
                logger.error(f"Alert callback error: {e}")
    
    async def log_auth_attempt(self, device_id: str, success: bool, 
                             session_id: str = None, details: Dict = None):
        """Log authentication attempt."""
        event = SecurityEvent(
            event_id=self._generate_event_id(),
            event_type=EventType.AUTH_SUCCESS if success else EventType.AUTH_FAILURE,
            severity=Severity.LOW if success else Severity.MEDIUM,
            device_id=device_id,
            session_id=session_id
        )
        
        event.add_detail("success", success)
        if details:
            for key, value in details.items():
                event.add_detail(key, value)
        
        event.add_tag("authentication")
        
        self._trigger_alerts(event)
        self.audit_logger.log_event(event)
    
    async def log_pairing_event(self, device_id: str, action: str, 
                              success: bool, details: Dict = None):
        """Log device pairing event."""
        if action == "request":
            event_type = EventType.PAIRING_REQUEST
        elif success:
            event_type = EventType.PAIRING_SUCCESS
        else:
            event_type = EventType.PAIRING_FAILURE
        
        event = SecurityEvent(
            event_id=self._generate_event_id(),
            event_type=event_type,
            severity=Severity.LOW if success else Severity.MEDIUM,
            device_id=device_id
        )
        
        event.add_detail("action", action)
        event.add_detail("success", success)
        if details:
            for key, value in details.items():
                event.add_detail(key, value)
        
        event.add_tag("pairing")
        
        self.audit_logger.log_event(event)
    
    async def log_session_event(self, event_type: EventType, session_id: str,
                              device_id: str = None, details: Dict = None):
        """Log session management event."""
        event = SecurityEvent(
            event_id=self._generate_event_id(),
            event_type=event_type,
            severity=Severity.LOW,
            session_id=session_id,
            device_id=device_id
        )
        
        if details:
            for key, value in details.items():
                event.add_detail(key, value)
        
        event.add_tag("session")
        
        self.audit_logger.log_event(event)
    
    async def log_key_operation(self, operation: str, key_id: str,
                              session_id: str = None, details: Dict = None):
        """Log cryptographic key operation."""
        event_type_map = {
            "generate": EventType.KEY_GENERATED,
            "rotate": EventType.KEY_ROTATED,
            "revoke": EventType.KEY_REVOKED
        }
        
        event = SecurityEvent(
            event_id=self._generate_event_id(),
            event_type=event_type_map.get(operation, EventType.KEY_GENERATED),
            severity=Severity.LOW,
            session_id=session_id
        )
        
        event.add_detail("operation", operation)
        event.add_detail("key_id", key_id)
        if details:
            for key, value in details.items():
                event.add_detail(key, value)
        
        event.add_tag("cryptography")
        
        self.audit_logger.log_event(event)
    
    async def log_security_violation(self, violation_type: str, device_id: str,
                                   severity: Severity = Severity.HIGH,
                                   details: Dict = None):
        """Log security violation."""
        event_type_map = {
            "replay_attack": EventType.REPLAY_ATTACK,
            "invalid_message": EventType.INVALID_MESSAGE,
            "device_blocked": EventType.DEVICE_BLOCKED,
            "device_unblocked": EventType.DEVICE_UNBLOCKED
        }
        
        event = SecurityEvent(
            event_id=self._generate_event_id(),
            event_type=event_type_map.get(violation_type, EventType.SECURITY_VIOLATION),
            severity=severity,
            device_id=device_id
        )
        
        event.add_detail("violation_type", violation_type)
        if details:
            for key, value in details.items():
                event.add_detail(key, value)
        
        event.add_tag("security_violation")
        
        self._trigger_alerts(event)
        self.audit_logger.log_event(event)
    
    async def log_system_event(self, event_name: str, severity: Severity,
                             details: Dict = None):
        """Log system-level event."""
        event_type_map = {
            "system_start": EventType.SYSTEM_START,
            "system_stop": EventType.SYSTEM_STOP
        }
        
        event = SecurityEvent(
            event_id=self._generate_event_id(),
            event_type=event_type_map.get(event_name, EventType.SYSTEM_START),
            severity=severity
        )
        
        event.add_detail("event_name", event_name)
        if details:
            for key, value in details.items():
                event.add_detail(key, value)
        
        event.add_tag("system")
        
        self.audit_logger.log_event(event)
    
    async def get_security_report(self, hours: int = 24) -> Dict[str, Any]:
        """Generate security report for specified time period."""
        end_time = time.time()
        start_time = end_time - (hours * 3600)
        
        # Read events from logs
        events = await self.audit_logger.read_events(start_time, end_time)
        
        # Analyze events
        report = {
            "report_period": {
                "start_time": start_time,
                "end_time": end_time,
                "hours": hours
            },
            "summary": {
                "total_events": len(events),
                "unique_devices": len(set(e.device_id for e in events if e.device_id)),
                "unique_sessions": len(set(e.session_id for e in events if e.session_id))
            },
            "event_counts": defaultdict(int),
            "severity_counts": defaultdict(int),
            "device_activity": defaultdict(int),
            "hourly_distribution": defaultdict(int),
            "top_violations": []
        }
        
        # Count events by type and severity
        violations = []
        for event in events:
            report["event_counts"][event.event_type.value] += 1
            report["severity_counts"][event.severity.value] += 1
            
            if event.device_id:
                report["device_activity"][event.device_id] += 1
            
            # Hourly distribution
            hour = int(event.timestamp / 3600)
            report["hourly_distribution"][hour] += 1
            
            # Track violations
            if event.event_type in [EventType.SECURITY_VIOLATION, 
                                  EventType.REPLAY_ATTACK,
                                  EventType.AUTH_FAILURE]:
                violations.append(event)
        
        # Top violations
        violation_counts = defaultdict(int)
        for violation in violations:
            key = f"{violation.event_type.value}_{violation.device_id}"
            violation_counts[key] += 1
        
        report["top_violations"] = sorted(
            violation_counts.items(),
            key=lambda x: x[1],
            reverse=True
        )[:10]
        
        # Convert defaultdicts to regular dicts for JSON serialization
        for key in ["event_counts", "severity_counts", "device_activity", 
                   "hourly_distribution"]:
            report[key] = dict(report[key])
        
        return report
    
    async def cleanup(self):
        """Clean up security audit system."""
        await self.stop()
        logger.info("SecurityAudit cleanup completed")


# Example usage and testing
if __name__ == "__main__":
    import asyncio
    import tempfile
    
    async def test_security_audit():
        """Test security audit functionality."""
        # Create temporary directory for testing
        with tempfile.TemporaryDirectory() as temp_dir:
            audit = SecurityAudit(temp_dir)
            
            # Set up alert callback
            def security_alert(alert_type, details):
                print(f"🚨 SECURITY ALERT: {alert_type}")
                print(f"   Details: {details}")
            
            audit.add_alert_callback(security_alert)
            
            # Start audit system
            await audit.start()
            
            try:
                # Test various audit events
                print("=== Testing Security Audit Events ===")
                
                # Test authentication events
                await audit.log_auth_attempt("device_001", True, "session_123")
                await audit.log_auth_attempt("device_002", False, details={
                    "reason": "invalid_pin", "attempt": 1
                })
                
                # Test pairing events
                await audit.log_pairing_event("device_001", "request", True, {
                    "device_name": "Test Device", "trust_level": 2
                })
                
                # Test session events
                await audit.log_session_event(
                    EventType.SESSION_CREATED, "session_123", "device_001"
                )
                
                # Test key operations
                await audit.log_key_operation("generate", "key_456", "session_123")
                await audit.log_key_operation("rotate", "key_456", "session_123")
                
                # Test security violations
                await audit.log_security_violation(
                    "replay_attack", "device_002", Severity.HIGH, {
                        "sequence": 12345, "expected": 12346
                    }
                )
                
                await audit.log_security_violation(
                    "device_blocked", "device_002", Severity.MEDIUM, {
                        "reason": "too_many_failures", "duration": 300
                    }
                )
                
                # Wait for events to be flushed
                await asyncio.sleep(1)
                
                # Generate security report
                report = await audit.get_security_report(1)
                
                print(f"\n=== Security Report (1 hour) ===")
                print(f"Total events: {report['summary']['total_events']}")
                print(f"Unique devices: {report['summary']['unique_devices']}")
                print(f"Event types: {report['event_counts']}")
                print(f"Severity levels: {report['severity_counts']}")
                print(f"Top violations: {report['top_violations']}")
                
            finally:
                # Cleanup
                await audit.cleanup()
    
    # Run test
    asyncio.run(test_security_audit())